import torch
import torch.nn as nn

from monai.metrics import DiceMetric, SurfaceDiceMetric, HausdorffDistanceMetric


class SegMetrics(nn.Module):
    def __init__(self, num_classes) -> None:
        super().__init__()
        
        self.dice = DiceMetric(include_background=True,
                               num_classes=num_classes)

        self.nsd = SurfaceDiceMetric(include_background=True,
                               class_thresholds=[4, 4, 2])

        self.hd = HausdorffDistanceMetric(include_background=True,
                               percentile=95)
        
    def forward(self, x, y):
        dice = self.dice(x, y)
        nsd = self.nsd(x, y)
        hd = torch.nan_to_num(self.hd(x, y), 0.)
        
        return dice, nsd, hd
    
class MyDice(nn.Module):
    def __init__(self, n_class):
        super().__init__()
        self.n_class = n_class

    def compute_stats_brats(self, p, y):
        scores = torch.zeros(self.n_class, device=p.device, dtype=torch.float32)
        p = (torch.sigmoid(p) > 0.5).int()

        for i in range(self.n_class):
            p_i, y_i = p[:, i], y[:, i]
            if (y_i != 1).all():
                # no foreground class
                scores[i - 1] += 1 if (p_i != 1).all() else 0
                continue
            tp, fn, fp = self.get_stats(p_i, y_i, 1)
            denom = (2 * tp + fp + fn).to(torch.float)
            score_cls = (2 * tp).to(torch.float) / denom if torch.is_nonzero(denom) else 0.0
            scores[i] += score_cls
        return scores

    @staticmethod
    def get_stats(p, y, c):
        tp = torch.logical_and(p == c, y == c).sum()
        fn = torch.logical_and(p != c, y == c).sum()
        fp = torch.logical_and(p == c, y != c).sum()
        return tp, fn, fp

    def forward(self, p, y):
        self.dice += self.compute_stats_brats(p, y) if self.brats else self.compute_stats(p, y)
        return self.dice